from PyQt6.QtWidgets import QTreeWidget, QTreeWidgetItem, QMenu, QMessageBox
from PyQt6.QtGui import QIcon, QAction
from PyQt6.QtCore import Qt, pyqtSignal
from sqlalchemy import inspect
from sqlalchemy.exc import SQLAlchemyError


class DatabaseTree(QTreeWidget):
    connection_selected = pyqtSignal(str)
    table_selected = pyqtSignal(str, str)

    def __init__(self, connection_manager=None):
        super().__init__()
        self.connection_manager = connection_manager
        self.setHeaderLabel("数据库")
        self.setColumnCount(1)
        self.setAnimated(True)
        self.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.customContextMenuRequested.connect(self.show_context_menu)
        self.itemDoubleClicked.connect(self.on_item_double_clicked)

        self.populate_tree()

    # ----------------------------------------------------------
    # 真实连接树
    # ----------------------------------------------------------
    def populate_tree(self):
        """填充真实连接"""
        self.clear()

        if not self.connection_manager:
            return

        for conn_id, conn_info in self.connection_manager.connections.items():
            conn_item = QTreeWidgetItem([conn_info["name"]])
            conn_item.setIcon(0, QIcon("resources/icons/server.png"))
            conn_item.setData(0, Qt.ItemDataRole.UserRole,
                              {"type": "connection", "id": conn_id})
            self.addTopLevelItem(conn_item)

            # 占位符，点击后真正加载
            placeholder = QTreeWidgetItem(["双击连接"])
            conn_item.addChild(placeholder)

    # ----------------------------------------------------------
    # 连接服务器 -> 加载数据库
    # ----------------------------------------------------------
    def connect_to_server(self, conn_item):
        conn_id = conn_item.data(0, Qt.ItemDataRole.UserRole)["id"]
        conn_info = self.connection_manager.connections[conn_id]

        try:
            engine = self.connection_manager.get_engine(conn_id)
            if not engine:
                raise RuntimeError("无法获取数据库引擎")

            inspector = inspect(engine)
            conn_item.takeChildren()  # 移除占位

            for db_name in inspector.get_schema_names():
                db_item = QTreeWidgetItem([db_name])
                db_item.setIcon(0, QIcon("resources/icons/database.png"))
                db_item.setData(0, Qt.ItemDataRole.UserRole,
                                {"type": "database",
                                 "connection_id": conn_id,
                                 "schema": db_name})
                conn_item.addChild(db_item)

                # 同样延迟加载表
                placeholder = QTreeWidgetItem(["双击加载表"])
                db_item.addChild(placeholder)

            conn_item.setExpanded(True)
            self.connection_selected.emit(conn_id)
            self.connection_manager.current_connection = conn_id   # ➕ 新增
        except SQLAlchemyError as e:
            QMessageBox.critical(self, "连接失败", str(e))

    # ----------------------------------------------------------
    # 加载数据库下的表
    # ----------------------------------------------------------
    def refresh_database(self, db_item):
        conn_id = db_item.data(0, Qt.ItemDataRole.UserRole)["connection_id"]
        schema = db_item.data(0, Qt.ItemDataRole.UserRole)["schema"]

        try:
            engine = self.connection_manager.get_engine(conn_id)
            inspector = inspect(engine)

            db_item.takeChildren()
            for table_name in inspector.get_table_names(schema=schema):
                table_item = QTreeWidgetItem([table_name])
                table_item.setIcon(0, QIcon("resources/icons/table.png"))
                table_item.setData(0, Qt.ItemDataRole.UserRole,
                                   {"type": "table",
                                    "connection_id": conn_id,
                                    "schema": schema,
                                    "table": table_name})
                db_item.addChild(table_item)

            db_item.setExpanded(True)
        except SQLAlchemyError as e:
            QMessageBox.critical(self, "加载表失败", str(e))

    # ----------------------------------------------------------
    # 双击事件
    # ----------------------------------------------------------
    def on_item_double_clicked(self, item, _):
        data = item.data(0, Qt.ItemDataRole.UserRole)
        if not data:
            return

        if data["type"] == "connection":
            self.connect_to_server(item)
        elif data["type"] == "database":
            self.refresh_database(item)
        elif data["type"] == "table":
            schema = data["schema"]
            table = data["table"]
            self.table_selected.emit(schema, table)

    # ----------------------------------------------------------
    # 右键菜单（简化示例）
    # ----------------------------------------------------------
    def show_context_menu(self, pos):
        item = self.itemAt(pos)
        if not item:
            return

        data = item.data(0, Qt.ItemDataRole.UserRole)
        menu = QMenu()

        if data["type"] == "connection":
            act = QAction("连接", self)
            act.triggered.connect(lambda: self.connect_to_server(item))
            menu.addAction(act)

        elif data["type"] == "database":
            act = QAction("刷新", self)
            act.triggered.connect(lambda: self.refresh_database(item))
            menu.addAction(act)

        elif data["type"] == "table":
            act = QAction("SELECT * FROM", self)
            act.triggered.connect(lambda: self.table_selected.emit(
                data["schema"], data["table"]))
            menu.addAction(act)

        menu.exec(self.viewport().mapToGlobal(pos))